# funcomp.py --- function composition in python
#
# Copyright 2010 Martin Wiebusch
#
# This module is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this module.  If not, see <http://www.gnu.org/licenses/>.


import inspect

def _unpack(f, args):
    if isinstance(args, tuple):
        return f(*args)
    else:
        return f(args)

class ComposableFunction:
    _func = None

    def __init__(self, f):
        if isinstance(f, ComposableFunction):
            self._func = f._func
        elif inspect.isroutine(f) or inspect.isclass(f):
            self._func = f
        else:
            self._func = (lambda *args: f)

    # string representations
    def __repr__(self):
        return "funcomp.ComposableFunction("+self._func.__repr__()+")"
    def __str__(self):
        return self._func.__str__()

    # calls
    def __call__(self, *args):
        return self._func(*args)

    # composition
    def __mod__(self, f):
        f = ComposableFunction(f)._func
        return ComposableFunction(lambda *args: _unpack(self._func, f(*args)))
    def __rmod__(self, f):
        f = ComposableFunction(f)._func
        return ComposableFunction(lambda *args: _unpack(f, self._func(*args)))
    def __floordiv__(self, f):
        f = ComposableFunction(f)._func
        return ComposableFunction(lambda *args: self._func(f(*args)))
    def __rfloordiv__(self, f):
        f = ComposableFunction(f)._func
        return ComposableFunction(lambda *args: f(self._func(*args)))
    def __xor__(self, n):
        if isinstance(n, int) and n > 0:
            f = self
            for i in range(n-1):
                f = self % f
            return f
        elif n == 0:
            return identity
        else:
            return NotImplemented
            
    # binary operations
    def __add__(self, f):
        f = ComposableFunction(f)._func
        return ComposableFunction(lambda *args: self._func(*args) + f(*args))
    def __radd__(self, f):
        f = ComposableFunction(f)._func
        return ComposableFunction(lambda *args: f(*args) + self._func(*args))
    def __sub__(self, f):
        f = ComposableFunction(f)._func
        return ComposableFunction(lambda *args: self._func(*args) - f(*args))
    def __rsub__(self, f):
        f = ComposableFunction(f)._func
        return ComposableFunction(lambda *args: f(*args) - self._func(*args))
    def __mul__(self, f):
        f = ComposableFunction(f)._func
        return ComposableFunction(lambda *args: self._func(*args) * f(*args))
    def __rmul__(self, f):
        f = ComposableFunction(f)._func
        return ComposableFunction(lambda *args: f(*args) * self._func(*args))
    def __div__(self, f):
        f = ComposableFunction(f)._func
        return ComposableFunction(lambda *args: self._func(*args) / f(*args))
    def __rdiv__(self, f):
        f = ComposableFunction(f)._func
        return ComposableFunction(lambda *args: f(*args) / self._func(*args))

    # unary operations
    def __pow__(self, n):
        return ComposableFunction(lambda *args: self._func(*args)**n)
    def __neg__(self):
        return ComposableFunction(lambda *args: -self._func(*args))
    def __pos__(self):
        return self
    def __abs__(self):
        return ComposableFunction(lambda *args: abs(self._func(*args)))
    
    # comparisons
    def __lt__(self, f):
        f = ComposableFunction(f)._func
        return ComposableFunction(lambda *args: self._func(*args) < f(*args))
    def __le__(self, f):
        f = ComposableFunction(f)._func
        return ComposableFunction(lambda *args: self._func(*args) <= f(*args))
    def __eq__(self, f):
        f = ComposableFunction(f)._func
        return ComposableFunction(lambda *args: self._func(*args) == f(*args))
    def __ne__(self, f):
        f = ComposableFunction(f)._func
        return ComposableFunction(lambda *args: self._func(*args) != f(*args))
    def __ge__(self, f):
        f = ComposableFunction(f)._func
        return ComposableFunction(lambda *args: self._func(*args) >= f(*args))
    def __gt__(self, f):
        f = ComposableFunction(f)._func
        return ComposableFunction(lambda *args: self._func(*args) > f(*args))

    # logical operations
    def __and__(self, f):
        f = ComposableFunction(f)._func
        return ComposableFunction(lambda *args: self._func(*args) and f(*args))
    def __rand__(self, f):
        f = ComposableFunction(f)._func
        return ComposableFunction(lambda *args: f(*args) and self._func(*args))
    def __or__(self, f):
        f = ComposableFunction(f)._func
        return ComposableFunction(lambda *args: self._func(*args) or f(*args))
    def __ror__(self, f):
        f = ComposableFunction(f)._func
        return ComposableFunction(lambda *args: f(*args) or self._func(*args))
    def __invert__(self):
        return ComposableFunction(lambda *args: not self._func(*args))


@ComposableFunction
def identity(x):
    return x
